"""工作流引擎核心模块 - 提供可扩展的工作流管理和执行框架"""

from typing import List, Dict, Any, Optional, Callable, Type, Set, Union
import time
import importlib
import inspect
import os
from abc import ABC, abstractmethod
from concurrent.futures import ThreadPoolExecutor, as_completed
from enum import Enum

from rpa_framework.core.task import Task, TaskStatus
from rpa_framework.rpa_logging.logger import Logger
from rpa_framework.exceptions.rpa_exceptions import WorkflowException as WorkflowExecutionException


class NodeType(Enum):
    """节点类型枚举"""
    START = "start"
    END = "end"
    FILE = "file"
    BROWSER = "browser"
    DESKTOP = "desktop"
    DATA = "data"
    CONDITION = "condition"
    LOOP = "loop"
    WAIT = "wait"
    NOTIFICATION = "notification"


class NodeExecutor(ABC):
    """节点执行器抽象基类 - 所有节点类型的执行逻辑基类"""
    
    def __init__(self):
        self.context = None
    
    def execute(self, params: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
        """执行节点逻辑
        
        Args:
            params: 节点参数
            context: 执行上下文，可用于在节点间传递数据
            
        Returns:
            执行结果
        """
        # 保存上下文引用
        self.context = context
        
        # 解析参数中的变量引用
        resolved_params = self.resolve_variables(params)
        
        # 调用具体实现的执行逻辑
        return self._execute(resolved_params, context)
    
    @abstractmethod
    def _execute(self, params: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
        """具体的节点执行逻辑，由子类实现
        
        Args:
            params: 已解析变量引用的节点参数
            context: 执行上下文
            
        Returns:
            执行结果
        """
        pass
    
    def resolve_variables(self, params: Dict[str, Any]) -> Dict[str, Any]:
        """解析参数中的变量引用
        
        Args:
            params: 原始参数
            
        Returns:
            解析变量引用后的参数
        """
        import re
        
        def _resolve_value(value):
            if isinstance(value, str):
                # 支持 {{variable_name}} 和 ${variable_name} 两种格式
                # 先处理 {{variable_name}} 格式
                def replace_brackets(match):
                    var_name = match.group(1)
                    return str(self.get_variable(var_name, f"{{{{{var_name}}}}}"))
                
                value = re.sub(r'\{\{(.*?)\}\}', replace_brackets, value)
                
                # 再处理 ${variable_name} 格式
                def replace_dollar(match):
                    var_name = match.group(1)
                    return str(self.get_variable(var_name, f"${{{var_name}}}"))
                
                value = re.sub(r'\$\{(.*?)\}', replace_dollar, value)
            elif isinstance(value, dict):
                # 递归解析字典中的所有值
                for k, v in value.items():
                    value[k] = _resolve_value(v)
            elif isinstance(value, list):
                # 递归解析列表中的所有值
                for i in range(len(value)):
                    value[i] = _resolve_value(value[i])
            return value
        
        # 深拷贝参数以避免修改原始参数
        import copy
        resolved = copy.deepcopy(params)
        
        # 解析所有参数值中的变量引用
        for key, value in resolved.items():
            resolved[key] = _resolve_value(value)
        
        return resolved
    
    def get_variable(self, name: str, default: Any = None) -> Any:
        """从上下文中获取变量
        
        Args:
            name: 变量名
            default: 默认值
            
        Returns:
            变量值或默认值
        """
        if not self.context:
            return default
            
        # 首先尝试从context中直接获取（通过set_variable设置的值）
        if isinstance(self.context, dict):
            if 'variables' in self.context and name in self.context['variables']:
                return self.context['variables'][name]
            # 然后尝试从普通上下文中获取
            return self.context.get(name, default)
        
        # 如果context是WorkflowExecutionContext实例
        elif hasattr(self.context, 'get_variable'):
            return self.context.get_variable(name, default)
        elif hasattr(self.context, 'get'):
            return self.context.get(name, default)
        
        return default
    
    def set_variable(self, name: str, value: Any):
        """设置变量到上下文中
        
        Args:
            name: 变量名
            value: 变量值
        """
        if not self.context:
            return
            
        if isinstance(self.context, dict):
            # 确保variables字典存在
            if 'variables' not in self.context:
                self.context['variables'] = {}
            self.context['variables'][name] = value
            # 同时也设置到普通上下文中，便于直接访问
            self.context[name] = value
        elif hasattr(self.context, 'set_variable'):
            self.context.set_variable(name, value)
        elif hasattr(self.context, 'set'):
            self.context.set(name, value)
    
    @property
    def node_type(self) -> str:
        """获取节点类型"""
        return self.__class__.__name__.replace("Executor", "").lower()


class StartNodeExecutor(NodeExecutor):
    """开始节点执行器"""
    
    def _execute(self, params: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
        """执行开始节点"""
        context['workflow_start_time'] = time.time()
        
        # 初始化工作流变量（如果有）
        if 'variables' in params:
            for var_name, var_value in params['variables'].items():
                self.set_variable(var_name, var_value)
        
        return {
            "status": "success", 
            "message": "工作流开始执行",
            "data": {"start_time": context['workflow_start_time']}
        }


class EndNodeExecutor(NodeExecutor):
    """结束节点执行器"""
    
    def _execute(self, params: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
        """执行结束节点"""
        end_time = time.time()
        start_time = context.get('workflow_start_time', end_time)
        execution_time = end_time - start_time
        
        # 收集所有变量作为最终结果的一部分
        variables = {}
        if isinstance(context, dict):
            variables = context.get('variables', {})
        elif hasattr(context, 'variables'):
            variables = context.variables
        
        return {
            "status": "success",
            "message": "工作流执行结束",
            "data": {
                "end_time": end_time,
                "execution_time": execution_time,
                "variables": variables
            }
        }


class NodeRegistry:
    """节点注册表 - 管理所有可用的节点类型及其执行器"""
    
    _instance = None
    
    def __new__(cls):
        """单例模式"""
        if cls._instance is None:
            cls._instance = super(NodeRegistry, cls).__new__(cls)
            cls._instance._executors: Dict[str, Type[NodeExecutor]] = {}
            cls._instance._register_builtin_nodes()
        return cls._instance
    
    def _register_builtin_nodes(self):
        """注册内置节点类型"""
        self.register_executor(StartNodeExecutor)
        self.register_executor(EndNodeExecutor)
    
    def register_executor(self, executor_class: Type[NodeExecutor]):
        """注册节点执行器
        
        Args:
            executor_class: 节点执行器类
        """
        executor = executor_class()
        node_type = executor.node_type
        self._executors[node_type] = executor_class
        Logger().debug(f"注册节点执行器: {node_type} -> {executor_class.__name__}")
    
    def get_executor(self, node_type: str) -> NodeExecutor:
        """获取节点执行器实例
        
        Args:
            node_type: 节点类型
            
        Returns:
            节点执行器实例
            
        Raises:
            ValueError: 如果节点类型未注册
        """
        if node_type not in self._executors:
            raise ValueError(f"未知的节点类型: {node_type}")
        return self._executors[node_type]()
    
    def list_available_nodes(self) -> List[str]:
        """列出所有可用的节点类型
        
        Returns:
            节点类型列表
        """
        return list(self._executors.keys())
    
    def load_plugins(self, plugin_dir: str):
        """从插件目录加载节点插件
        
        Args:
            plugin_dir: 插件目录路径
        """
        if not os.path.exists(plugin_dir):
            Logger().warning(f"插件目录不存在: {plugin_dir}")
            return
        
        for filename in os.listdir(plugin_dir):
            if filename.endswith('.py') and not filename.startswith('_'):
                module_name = filename[:-3]
                try:
                    # 动态导入模块
                    spec = importlib.util.spec_from_file_location(module_name, os.path.join(plugin_dir, filename))
                    if spec and spec.loader:
                        module = importlib.util.module_from_spec(spec)
                        spec.loader.exec_module(module)
                        
                        # 查找并注册所有NodeExecutor子类
                        for name, obj in inspect.getmembers(module):
                            if (inspect.isclass(obj) and 
                                issubclass(obj, NodeExecutor) and 
                                obj is not NodeExecutor and 
                                not inspect.isabstract(obj)):
                                self.register_executor(obj)
                except Exception as e:
                    Logger().error(f"加载插件 {filename} 失败: {str(e)}")


class WorkflowExecutionContext:
    """工作流执行上下文 - 用于在工作流执行过程中共享数据"""
    
    def __init__(self):
        self.data = {}
        self.task_results = {}
        self.variables = {}
    
    def set(self, key: str, value: Any):
        """设置上下文数据"""
        self.data[key] = value
    
    def get(self, key: str, default: Any = None) -> Any:
        """获取上下文数据"""
        return self.data.get(key, default)
    
    def store_task_result(self, task_id: str, result: Any):
        """存储任务执行结果"""
        self.task_results[task_id] = result
    
    def get_task_result(self, task_id: str, default: Any = None) -> Any:
        """获取任务执行结果"""
        return self.task_results.get(task_id, default)
    
    def set_variable(self, name: str, value: Any):
        """设置工作流变量 - 用于在节点间传递数据"""
        self.variables[name] = value
    
    def get_variable(self, name: str, default: Any = None) -> Any:
        """获取工作流变量"""
        return self.variables.get(name, default)
    
    def resolve_variable_references(self, value: Any) -> Any:
        """解析值中的变量引用
        
        支持格式: {{variable_name}} 或 ${variable_name}
        
        Args:
            value: 可能包含变量引用的值
            
        Returns:
            解析后的值
        """
        if isinstance(value, str):
            import re
            # 匹配 {{variable_name}} 格式
            def replace_brackets(match):
                var_name = match.group(1)
                return str(self.get_variable(var_name, f"{{{{{var_name}}}}}"))
            
            value = re.sub(r'\{\{(.*?)\}\}', replace_brackets, value)
            
            # 匹配 ${variable_name} 格式
            def replace_dollar(match):
                var_name = match.group(1)
                return str(self.get_variable(var_name, f"${{{var_name}}}"))
            
            value = re.sub(r'\$\{(.*?)\}', replace_dollar, value)
            
        elif isinstance(value, dict):
            # 递归解析字典中的所有值
            for k, v in value.items():
                value[k] = self.resolve_variable_references(v)
        elif isinstance(value, list):
            # 递归解析列表中的所有值
            for i in range(len(value)):
                value[i] = self.resolve_variable_references(value[i])
        
        return value


class Workflow:
    """工作流类 - 支持可扩展的工作流管理和执行"""
    
    def __init__(self, name: str, description: str = ""):
        """初始化工作流
        
        Args:
            name: 工作流名称
            description: 工作流描述
        """
        self.name = name
        self.description = description
        self.tasks: List[Task] = []
        self.logger = Logger(f"Workflow-{name}")
        self.start_time = None
        self.end_time = None
        self.node_registry = NodeRegistry()
    
    def add_task(self, task: Task) -> 'Workflow':
        """添加任务到工作流
        
        Args:
            task: 要添加的任务
            
        Returns:
            工作流实例，支持链式调用
        """
        self.tasks.append(task)
        self.logger.debug(f"添加任务到工作流: {task.name} (ID: {task.id})")
        return self
    
    def create_task(self, 
                    name: str,
                    action = None,
                    node_type: str = None,
                    node_params: Dict[str, Any] = None,
                    args: List[Any] = None,
                    kwargs: Dict[str, Any] = None,
                    dependencies: List[Task] = None,
                    timeout: int = None,
                    max_retries: int = 0,
                    retry_delay: int = 1,
                    description: str = "") -> Task:
        """在工作流中创建新任务
        
        Args:
            name: 任务名称
            action: 要执行的函数（传统方式）
            node_type: 节点类型（如果使用节点注册机制）
            node_params: 节点参数（如果使用节点注册机制）
            args: 位置参数列表
            kwargs: 关键字参数字典
            dependencies: 依赖的其他任务列表
            timeout: 任务超时时间（秒）
            max_retries: 最大重试次数
            retry_delay: 重试延迟时间（秒）
            description: 任务描述
            
        Returns:
            创建的任务实例
        """
        # 创建任务对象
        task = Task(
            name=name,
            action=None,  # 稍后设置
            args=args or [],
            kwargs=kwargs or {},
            dependencies=dependencies or [],
            timeout=timeout,
            description=description
        )
        
        # 添加重试信息
        task.max_retries = max_retries
        task.retry_delay = retry_delay
        
        # 如果指定了节点类型，使用节点执行器
        if node_type:
            # 将节点信息存储在任务对象的自定义属性中
            task.node_type = node_type
            task.node_params = node_params
            
            # 创建包装函数，使用节点执行器执行
            # 注意：这里使用闭包，确保node_wrapper能访问到正确的task和全局context
            def create_node_wrapper(current_task):
                def node_wrapper():
                    try:
                        executor = self.node_registry.get_executor(node_type)
                        params = node_params or {}
                        
                        # 为每个任务提供对全局上下文的访问
                        # 如果没有上下文，则创建一个新的
                        if not hasattr(current_task, 'execution_context'):
                            # 确保任务能够访问到工作流的执行上下文
                            if context and hasattr(self, '_current_context'):
                                current_task.execution_context = self._current_context
                            else:
                                current_task.execution_context = WorkflowExecutionContext()
                        
                        # 执行节点逻辑，传入参数和上下文数据
                        result = executor.execute(params, current_task.execution_context.data)
                        
                        # 将任务结果存储在上下文中，以便其他任务可以访问
                        if self._current_context:
                            self._current_context.store_task_result(current_task.id, result)
                        
                        return result
                    except Exception as e:
                        self.logger.error(f"节点执行失败: {str(e)}")
                        raise
                return node_wrapper
            
            task_action = create_node_wrapper(task)
        else:
            # 使用传统的action方式
            task_action = action
        
        # 设置任务的action
        task.action = task_action
        
        self.add_task(task)
        return task
    
    def get_task(self, task_id: Union[str, None] = None, task_name: Union[str, None] = None) -> Optional[Task]:
        """根据ID或名称获取任务
        
        Args:
            task_id: 任务ID
            task_name: 任务名称
            
        Returns:
            找到的任务实例，如果没有找到返回None
        """
        for task in self.tasks:
            if task_id and task.id == task_id:
                return task
            if task_name and task.name == task_name:
                return task
        return None
    
    def validate(self) -> bool:
        """验证工作流的有效性
        
        Returns:
            工作流有效返回True，否则返回False
        """
        # 检查是否有循环依赖
        visited = set()
        path = set()
        
        def has_cycle(task, visited, path):
            visited.add(task)
            path.add(task)
            
            for dep in task.dependencies:
                if dep not in visited:
                    if has_cycle(dep, visited, path):
                        return True
                elif dep in path:
                    self.logger.error(f"检测到循环依赖: {dep.name} -> {task.name}")
                    return True
            
            path.remove(task)
            return False
        
        # 检查每个任务是否有循环依赖
        for task in self.tasks:
            if task not in visited:
                if has_cycle(task, visited, path):
                    return False
        
        # 检查任务ID唯一性
        task_ids = [task.id for task in self.tasks]
        if len(task_ids) != len(set(task_ids)):
            self.logger.error("工作流中存在重复的任务ID")
            return False
        
        # 检查节点类型是否有效
        for task in self.tasks:
            if hasattr(task, 'node_type'):
                node_type = task.node_type
                try:
                    self.node_registry.get_executor(node_type)
                except ValueError:
                    self.logger.error(f"无效的节点类型: {node_type} in task {task.name}")
                    return False
        
        return True
    
    def _topological_sort(self) -> List[Task]:
        """对任务进行拓扑排序
        
        Returns:
            排序后的任务列表
        """
        visited = {task: False for task in self.tasks}
        result = []
        
        def visit(task):
            if not visited[task]:
                visited[task] = True
                # 先访问所有依赖
                for dep in task.dependencies:
                    if dep in visited:  # 确保依赖任务在工作流中
                        visit(dep)
                # 然后将当前任务加入结果
                result.append(task)
        
        # 对每个未访问的任务进行访问
        for task in self.tasks:
            visit(task)
        
        return result
    
    def execute(self, 
                continue_on_error: bool = False,
                max_workers: int = 1,
                context: WorkflowExecutionContext = None) -> Dict[str, Any]:
        """执行工作流
        
        Args:
            continue_on_error: 发生错误时是否继续执行
            max_workers: 最大并行执行的工作线程数
            context: 执行上下文，如果为None则创建新的
            
        Returns:
            工作流执行结果
        """
        # 验证工作流
        if not self.validate():
            raise WorkflowExecutionException("工作流验证失败")
        
        # 初始化执行上下文
        if context is None:
            context = WorkflowExecutionContext()
        
        # 保存当前上下文到工作流实例，供节点执行时访问
        self._current_context = context
        
        self.logger.info(f"开始执行工作流: {self.name}")
        self.start_time = time.time()
        context.set('workflow_name', self.name)
        context.set('workflow_start_time', self.start_time)
        
        # 按依赖关系排序任务
        sorted_tasks = self._topological_sort()
        
        # 执行任务
        results = {}
        failed_tasks = []
        executed_tasks = 0
        
        # 如果设置了并行执行
        if max_workers > 1:
            # 分组执行：每组包含相互独立的任务
            task_groups = self._group_independent_tasks(sorted_tasks)
            
            for group in task_groups:
                with ThreadPoolExecutor(max_workers=min(max_workers, len(group))) as executor:
                    future_to_task = {executor.submit(self._execute_task_with_retry, task, context): task for task in group}
                    
                    for future in as_completed(future_to_task):
                        task = future_to_task[future]
                        try:
                            result = future.result()
                            results[task.name] = result
                            context.store_task_result(task.id, result)
                            executed_tasks += 1
                            self.logger.info(f"任务 {task.name} 执行成功")
                        except Exception as e:
                            failed_tasks.append(task)
                            self.logger.error(f"任务 {task.name} 执行失败: {str(e)}")
                            
                            if not continue_on_error:
                                self.end_time = time.time()
                                raise WorkflowExecutionException(
                                    f"工作流执行失败: 任务 {task.name} 执行出错") from e
        else:
            # 串行执行
            for task in sorted_tasks:
                try:
                    self.logger.info(f"执行任务: {task.name} (ID: {task.id})")
                    result = self._execute_task_with_retry(task, context)
                    results[task.name] = result
                    context.store_task_result(task.id, result)
                    executed_tasks += 1
                    self.logger.info(f"任务 {task.name} 执行成功")
                except Exception as e:
                    failed_tasks.append(task)
                    self.logger.error(f"任务 {task.name} 执行失败: {str(e)}")
                    
                    if not continue_on_error:
                        self.end_time = time.time()
                        raise WorkflowExecutionException(
                            f"工作流执行失败: 任务 {task.name} 执行出错") from e
        
        # 清理当前上下文引用
        self._current_context = None
        
        self.end_time = time.time()
        context.set('workflow_end_time', self.end_time)
        
        # 生成执行摘要
        summary = {
            "工作流名称": self.name,
            "总任务数": len(self.tasks),
            "已执行任务数": executed_tasks,
            "成功任务数": len(self.tasks) - len(failed_tasks),
            "失败任务数": len(failed_tasks),
            "执行时长": f"{(self.end_time - self.start_time):.2f}秒",
            "失败任务": [task.name for task in failed_tasks]
        }
        
        self.logger.info(f"工作流执行完成: {summary}")
        
        return {
            "summary": summary,
            "results": results,
            "tasks": [task.to_dict() for task in self.tasks],
            "context": context.data
        }
    
    def _execute_task_with_retry(self, task: Task, context: WorkflowExecutionContext) -> Any:
        """带重试机制的任务执行
        
        Args:
            task: 要执行的任务
            context: 执行上下文
            
        Returns:
            任务执行结果
        """
        max_retries = getattr(task, 'max_retries', 0)
        retry_delay = getattr(task, 'retry_delay', 1)
        
        # 为任务设置执行上下文
        task.execution_context = context
        
        for attempt in range(max_retries + 1):
            try:
                # 对于节点类型任务，我们需要特殊处理
                # 由于Task.execute()不接受额外参数，我们修改node_wrapper函数的实现方式
                # 这里直接执行任务，任务的action已经封装了context的处理
                result = task.execute()
                
                # 将结果存储在上下文中
                context.store_task_result(task.id, result)
                
                # 如果节点返回了数据，将其添加到上下文中
                # 这样其他节点可以通过引用前一个节点的输出
                if isinstance(result, dict) and 'data' in result:
                    if 'output_key' in getattr(task, 'node_params', {}):
                        # 使用用户指定的输出键
                        output_key = task.node_params['output_key']
                        context.set(output_key, result['data'])
                    else:
                        # 默认使用任务名称作为键
                        context.set(task.name, result['data'])
                
                return result
            except Exception as e:
                if attempt < max_retries:
                    self.logger.warning(f"任务 {task.name} 执行失败 (尝试 {attempt + 1}/{max_retries + 1}), {retry_delay}秒后重试: {str(e)}")
                    time.sleep(retry_delay)
                else:
                    self.logger.error(f"任务 {task.name} 执行失败，已达到最大重试次数: {str(e)}")
                    raise
    
    def _group_independent_tasks(self, sorted_tasks: List[Task]) -> List[List[Task]]:
        """将任务分组，同一组内的任务相互独立，可以并行执行
        
        Args:
            sorted_tasks: 拓扑排序后的任务列表
            
        Returns:
            任务组列表
        """
        groups = []
        task_to_group = {}
        
        for task in sorted_tasks:
            # 找出所有依赖任务所在的最大组索引
            max_group_idx = -1
            for dep in task.dependencies:
                if dep in task_to_group:
                    max_group_idx = max(max_group_idx, task_to_group[dep])
            
            # 当前任务应该放在最大组索引+1的位置
            current_group_idx = max_group_idx + 1
            
            # 如果当前组索引超出了现有组数，创建新组
            while len(groups) <= current_group_idx:
                groups.append([])
            
            # 将任务添加到相应的组
            groups[current_group_idx].append(task)
            task_to_group[task] = current_group_idx
        
        return groups
    
    def get_status(self) -> Dict[str, Any]:
        """获取工作流状态
        
        Returns:
            工作流状态信息
        """
        status_counts = {status: 0 for status in TaskStatus}
        
        for task in self.tasks:
            status_counts[task.status] += 1
        
        return {
            "工作流名称": self.name,
            "任务总数": len(self.tasks),
            "任务状态统计": {status.value: count for status, count in status_counts.items()},
            "开始时间": self.start_time,
            "结束时间": self.end_time,
            "执行时长": f"{(self.end_time - self.start_time):.2f}秒" if self.start_time and self.end_time else None,
            "可用节点类型": self.node_registry.list_available_nodes()
        }
    
    def to_dict(self) -> Dict[str, Any]:
        """将工作流信息转换为字典
        
        Returns:
            工作流信息字典
        """
        return {
            "name": self.name,
            "description": self.description,
            "task_count": len(self.tasks),
            "tasks": [task.to_dict() for task in self.tasks],
            "status": self.get_status(),
            "node_types": self.node_registry.list_available_nodes()
        }
    
    def load_plugin_nodes(self, plugin_dir: str):
        """加载插件节点
        
        Args:
            plugin_dir: 插件目录路径
        """
        self.node_registry.load_plugins(plugin_dir)
        self.logger.info(f"已加载插件节点，当前可用节点类型: {', '.join(self.node_registry.list_available_nodes())}")


# 扩展Task类，添加重试相关属性
Task.max_retries = 0
Task.retry_delay = 1


if __name__ == "__main__":
    # 示例用法
    def task1():
        print("执行任务1")
        return "任务1结果"
    
    def task2():
        print("执行任务2")
        return "任务2结果"
    
    # 创建工作流
    workflow = Workflow("示例工作流", "这是一个示例工作流")
    
    # 传统方式创建任务
    task1_obj = workflow.create_task("任务1", task1)
    task2_obj = workflow.create_task("任务2", task2, dependencies=[task1_obj])
    
    # 使用节点类型创建任务
    workflow.create_task("开始节点", node_type="start")
    workflow.create_task("结束节点", node_type="end")
    
    # 执行工作流
    result = workflow.execute()
    print(result)